import numpy as np
import matplotlib.pyplot as plt

# https://www.youtube.com/watch?v=5w5iUbTlpMQ

class KMeansClustering:
    def __init__(self, k=3):
        self.k = k
        self.centroids = None
    
    @staticmethod
    def euclidean_distance(data_point, centroids):
        return np.sqrt(np.sum((centroids - data_point)**2, axis=1))
    
    def fit(self, X, max_iter=200):
        self.centroids = np.random.uniform(np.amin(X, axis=0), np.amax(X, axis=0),
                                           size=(self.k, X.shape[1]))
        for _ in range(max_iter):
            y = []
            for data_point in X:
                dis = KMeansClustering.euclidean_distance(data_point, self.centroids)
                cluster_num = np.argmin(dis)
                y.append(cluster_num)

            y = np.array(y)
            cluster_indices = []
            for i in range(self.k):
                cluster_indices.append(np.argwhere(y == i))

            cluster_centers = []
            for i, indices in enumerate(cluster_indices):
                if len(indices) == 0:
                    cluster_centers.append(self.centroids[i])
                else:
                    cluster_centers.append(np.mean(X[indices], axis=0)[0])
            
            if (np.max(self.centroids - np.array(cluster_centers)) < 0.0001):
                break
            else:
                self.centroids = np.array(cluster_centers)
        
        return y

from sklearn.datasets import make_blobs
from sklearn import metrics

# 生成各向同性的高斯斑点以进行聚类, 总点数为100；如果是数组，则序列中的每个元素表示每个簇的样本数
X, y, centers = make_blobs(n_samples=100, n_features=2, centers=3, return_centers=True)
random_points = X  # (100,2) 的数据集
# random_points = np.random.randint(0, 100, (100, 2))

kmeans = KMeansClustering(k=3)
labels = kmeans.fit(random_points)

# ground-truth class
print('true class:', y)
print('predicted class:', labels)

# 用ari 校验此算法生成的簇与ground-truth 差别大不大
ari = metrics.adjusted_rand_score(y, labels)
print("adjusted_rand_score:", ari) # 越接近1，越相似

print('predicted centers:', kmeans.centroids)

plt.scatter(random_points[:, 0], random_points[:, 1], c=labels)
# 把质心的x维， y维抽取出来即可
plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], c = range(len(kmeans.centroids)),
            marker='*', s=200)
plt.show()